#!/usr/bin/env python3
# src/main.py 

import asyncio
import click
import logging
import sys
import signal
import threading
import time
import os
from pathlib import Path
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn
from rich.logging import RichHandler
from rich.panel import Panel
from rich.text import Text

# 导入自定义模块
from .config_manager import ConfigManager
from .ssh_connection import SSHConnectionManager
from .command_executor import CommandExecutor
from .result_collector import ResultCollector

# 配置Rich控制台
console = Console()

# 配置日志系统
logging.basicConfig(
    level=logging.INFO,
    format="%(message)s",
    datefmt="[%X]",
    handlers=[RichHandler(rich_tracebacks=True, console=console)]
)

logger = logging.getLogger(__name__)


def validate_config_file(ctx, param, value):
    """Click回调函数：验证配置文件是否存在"""
    if value is None:
        return value
    
    config_path = Path(value)
    if not config_path.exists():
        raise click.BadParameter(f"配置文件不存在: {value}")
    if not config_path.is_file():
        raise click.BadParameter(f"路径不是文件: {value}")
    
    return str(config_path.absolute())

class GracefulShutdown:
    """优雅关闭管理器"""
    
    def __init__(self):
        self.shutdown_event = None  # 将在asyncio环境中初始化
        self.shutdown_requested = False
        self.active_tasks = set()
        self.ssh_manager = SSHConnectionManager()
        self.executor = CommandExecutor(self.ssh_manager)
        self.ctrl_c_count = 0
        self.last_ctrl_c_time = 0
        
    def init_async(self):
        """在asyncio环境中初始化"""
        if not self.shutdown_event:
            self.shutdown_event = asyncio.Event()
        
    def request_shutdown(self, force=False):
        """请求关闭"""
        current_time = time.time()
        
        # 检查是否为连续的Ctrl+C
        if current_time - self.last_ctrl_c_time < 2:
            self.ctrl_c_count += 1
        else:
            self.ctrl_c_count = 1
        
        self.last_ctrl_c_time = current_time
        
        if self.ctrl_c_count >= 2 or force:
            console.print("\n[red]强制退出！[/red]")
            os._exit(130)  # 强制退出
        
        if not self.shutdown_requested:
            self.shutdown_requested = True
            console.print("\n[yellow]正在关闭程序... (再次按Ctrl+C强制退出)[/yellow]")
            if self.shutdown_event:
                self.shutdown_event.set()
            
            # 设置一个超时强制退出
            async def force_exit():
                if self.shutdown_requested:
                    # 清理资源
                    asyncio.run(self.cleanup_resources(quiet=True))
                    # 等待所有任务完成
                    # await asyncio.sleep(0.5)
                    os._exit(130)
            
    async def cleanup_resources(self,quiet=False):
        """清理所有资源"""
        if not quiet:
            console.print("[dim]正在清理系统资源...[/dim]")
        
        
        # 清理执行器
        if self.executor:
            try:
                # 取消长时间运行的任务
                self.executor.cancel_all_tasks()
                await asyncio.sleep(0.5)

                if not quiet:
                    console.print("[dim]命令执行器已清理[/dim]")
            except Exception as e:
                console.print(f"[yellow]执行器清理异常: {str(e)}[/yellow]")


        # 清理SSH连接
        if self.ssh_manager:
            try:
                await asyncio.wait_for(
                    asyncio.get_event_loop().run_in_executor(
                        None, self.ssh_manager.close_all
                    ),
                    timeout=3.0
                )
                if not quiet:
                    console.print("[dim]SSH连接已关闭[/dim]")
            except asyncio.TimeoutError:
                console.print("[yellow]SSH连接关闭超时，强制终止[/yellow]")
            except Exception as e:
                console.print(f"[yellow]SSH连接关闭异常: {str(e)}[/yellow]")

# 全局关闭管理器
shutdown_manager = GracefulShutdown()

def setup_signal_handlers():
    """设置信号处理器（仅在主线程中调用）"""
    
    def signal_handler(signum, frame):
        """信号处理函数"""
        shutdown_manager.request_shutdown()
    
    # 只在主线程中设置信号处理器
    if threading.current_thread() is threading.main_thread():
        signal.signal(signal.SIGINT, signal_handler)
        signal.signal(signal.SIGTERM, signal_handler)
        logger.debug("信号处理器已设置")

async def run_validation(
    config_path: str, 
    report_name: str = '',
    quiet: bool = False,
    progress_callback=None
) -> int:
    """
    支持优雅关闭的验证执行函数
    """
    # 初始化异步环境
    shutdown_manager.init_async()
    
    config_manager = None
    ssh_manager = None
    executor = None
    collector = None
    
    try:
        # 第一步：加载配置
        if not quiet:
            console.print(Panel.fit(
                "[bold blue]正在加载配置文件...[/bold blue]",
                border_style="blue"
            ))
        
        # 检查是否已收到关闭信号
        if shutdown_manager.shutdown_requested:
            return 130
        
        config_manager = ConfigManager(config_path)
        
        if not config_manager.validate_config():
            console.print("[bold red]配置验证失败，程序终止[/bold red]")
            return 1
        
        stats = config_manager.get_statistics()
        if not quiet:
            console.print(f"✅ 配置加载成功: {stats['total_hosts']}台主机, {stats['total_tasks']}个任务")
        
        # 第二步：初始化组件
        if not quiet:
            console.print(Panel.fit(
                "[bold green]正在初始化系统组件...[/bold green]",
                border_style="green"
            ))
        
        if shutdown_manager.shutdown_requested:
            return 130
        
        ssh_manager = SSHConnectionManager()
        executor = CommandExecutor(ssh_manager)
        collector = ResultCollector()
        
        # 注册到关闭管理器
        shutdown_manager.ssh_manager = ssh_manager
        shutdown_manager.executor = executor
        
        # 第三步：执行验证任务（支持中断）
        if not quiet:
            console.print(Panel.fit(
                f"[bold yellow]开始执行验证任务...[/bold yellow]",
                border_style="yellow"
            ))

        logger.info("execute_tasks_with_interruption方法未找到，使用旧的execute_tasks方法")
        results = await execute_tasks_with_progress(
            executor, config_manager, progress_callback
        )
        
        # 生成报告
        if not quiet:
            console.print(Panel.fit(
                "[bold cyan]正在生成执行报告...[/bold cyan]",
                border_style="cyan"
            ))
        
        report_path = collector.generate_report(results, report_name)
        if not quiet:
            collector.print_summary(results)
            console.print(f"\n📄 详细报告已保存至: [bold green]{report_path}[/bold green]")

        # 可选：生成CSV摘要  
        try:  
            csv_path = collector.export_csv_summary(results)  
            if not quiet:
                console.print(f"📊 CSV摘要已保存至: [bold green]{csv_path}[/bold green]")  
        except Exception as e:  
            logger.warning(f"CSV报告生成失败: {str(e)}")  
        
        # 确定退出状态
        all_success = all(r.success for r in results)
        return 0 if all_success else 1
        
    except asyncio.CancelledError:
        console.print("[yellow]任务执行被取消[/yellow]")
        return 130
        
    except KeyboardInterrupt:
        console.print("\n[yellow]用户中断程序执行[/yellow]")
        return 130
        
    except Exception as e:
        console.print(f"[bold red]程序执行失败: {str(e)}[/bold red]")
        logger.exception("详细错误信息")
        return 1
        
    finally:
        # 执行清理
        await shutdown_manager.cleanup_resources(quiet)

async def execute_tasks_with_progress(executor, config_manager, progress_callback):
    """带进度显示的任务执行"""

    tasks = config_manager.tasks
    hosts = config_manager.hosts
    
    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        BarColumn(),
        TaskProgressColumn(),
        console=console,
        transient=False  # 保持进度条显示
    ) as progress:
        
        main_task = progress.add_task(
            "执行验证任务...", 
            total=len(tasks)
        )
        
        # 定义内部进度更新函数
        def update_progress(completed_tasks, total_tasks=None):
            """更新进度条"""
            if total_tasks is None:
                total_tasks = len(tasks)
            
            # 更新进度条
            progress.update(main_task, completed=completed_tasks)
            
            # 调用外部回调
            if progress_callback:
                logger.debug(f"调用外部进度回调: {completed_tasks}/{total_tasks}")
                progress_callback(completed_tasks, total_tasks)
            else:
                logger.debug(f"外部进度回调未设置")
            
            # 可选：显示当前进度信息
            percentage = (completed_tasks / total_tasks * 100) if total_tasks > 0 else 0
            logger.debug(f"任务进度: {completed_tasks}/{total_tasks} ({percentage:.1f}%)")

        # 创建一个监控关闭的任务
        async def monitor_shutdown():
            while not shutdown_manager.shutdown_requested:
                await asyncio.sleep(0.1)
            raise asyncio.CancelledError("Shutdown requested")
        
        shutdown_monitor = asyncio.create_task(monitor_shutdown())
        
        # 执行任务
        execution_task = asyncio.create_task(executor.execute_tasks(config_manager, progress_callback=update_progress))
        
        try:
            # 等待任务完成或关闭
            done, pending = await asyncio.wait(
                [execution_task, shutdown_monitor],
                return_when=asyncio.FIRST_COMPLETED
            )
            
            # 取消未完成的任务
            for task in pending:
                task.cancel()
            
            if execution_task in done:
                results = execution_task.result()
                progress.update(main_task, completed=len(tasks))
                return results
            else:
                # 被中断
                raise asyncio.CancelledError()
                
        except asyncio.CancelledError:
            console.print("[yellow]任务执行被中断[/yellow]")
            raise
        finally:
            # 清理监控任务
            if not shutdown_monitor.done():
                shutdown_monitor.cancel()


@click.command()
@click.option('--config', '-c', required=True, callback=validate_config_file,
              help='YAML配置文件路径', metavar='PATH')
@click.option('--report', '-r', help='自定义报告文件名（不含扩展名）', metavar='NAME')
@click.option('--log-level', '-l',
              type=click.Choice(['DEBUG', 'INFO', 'WARNING', 'ERROR'], case_sensitive=False),
              default='INFO', help='日志级别 [默认: INFO]')
@click.option('--quiet', '-q', is_flag=True, help='静默模式，只显示关键信息')
@click.option('--verbose', '-v', is_flag=True, help='详细模式，显示调试信息')
# 重复运行的次数，默认1
@click.option('--repeat-run',  default=1, help='重复运行的次数，默认1')
# 每次运行的间隔时间
@click.option('--interval',  default=0, help='每次运行的间隔时间，单位秒，默认0')
@click.option('--dry-run', is_flag=True, help='试运行模式，只解析配置不执行命令')
@click.version_option(version='1.0.0', message='%(prog)s %(version)s - Ubuntu Environment Validator')
def run(config, report, log_level, quiet, verbose, repeat_run, interval, dry_run):
    """Ubuntu环境自动化验证工具"""
    
    # 处理日志选项
    if quiet and verbose:
        console.print("[bold red]错误: --quiet 和 --verbose 选项不能同时使用[/bold red]")
        sys.exit(1)
    
    if verbose:
        log_level = 'DEBUG'
    elif quiet:
        log_level = 'WARNING'
    
    logging.getLogger().setLevel(getattr(logging, log_level.upper()))

    
    # 显示启动信息
    if not quiet:
        console.print(Panel.fit(
            Text("Ubuntu Environment Validator v1.0.0", style="bold blue"),
            subtitle="系统环境自动化验证工具",
            border_style="blue"
        ))
    
    # 设置信号处理器
    setup_signal_handlers()
    
    # 试运行模式
    if dry_run:
        try:
            console.print("[bold yellow]🔍 试运行模式 - 只验证配置文件[/bold yellow]")
            config_manager = ConfigManager(config)
            
            if config_manager.validate_config():
                stats = config_manager.get_statistics()
                console.print(Panel.fit(
                    f"[green]✅ 配置验证通过[/green]\n"
                    f"主机数量: {stats['total_hosts']}\n"
                    f"任务数量: {stats['total_tasks']}\n"
                    f"串行任务: {stats['serial_tasks']}\n"
                    f"并行任务: {stats['parallel_tasks']}",
                    title="配置统计",
                    border_style="green"
                ))
                sys.exit(0)
            else:
                console.print("[bold red]❌ 配置验证失败[/bold red]")
                sys.exit(1)
        except Exception as e:
            console.print(f"[bold red]配置验证异常: {str(e)}[/bold red]")
            sys.exit(1)
    
    # 正常执行模式
    exit_codes = []
    try:
        # 重复运行
        for i in range(repeat_run):
            console.print(f"\n[bold blue] 🔄 第 {i+1}/{repeat_run} 次运行[/bold blue]")
            # 运行异步验证流程
            exit_code = asyncio.run(run_validation(config, report, quiet))
            exit_codes.append(exit_code)

            if exit_code == 0:
                if not quiet:
                    console.print("\n[bold green]🎉 程序执行完成！[/bold green]")
            elif exit_code == 130:
                console.print("\n[yellow]⚠️  程序被用户中断[/yellow]")
                break
            else: 
                console.print(f"\n[bold red]❌ 第 {i+1} 次运行失败（退出码: {exit_code}）[/bold red]")
                console.print(f"\n[bold red]❌ 程序执行失败（退出码: {exit_code}）[/bold red]")
                break
        
            if interval > 0:
                console.print(f"\n[bold yellow]等待 {interval} 秒后进行下一次运行[/bold yellow]")
            time.sleep(interval)
        # 检查所有重复运行的退出码
        if all(code == 0 for code in exit_codes):
            console.print("\n[bold green]🎉 所有重复运行均成功！[/bold green]")
        else:
            console.print("\n[bold red]❌ 部分重复运行失败[/bold red]")
            output = '['+', '.join(f"{index+1}: {value}" for index, value in enumerate(exit_codes) if value != 0)+']'
            console.print(f"exit_codes:{output}")
        
    except KeyboardInterrupt:
        console.print("\n[yellow]程序被强制中断[/yellow]")
        sys.exit(130)
        
    except Exception as e:
        console.print(f"[bold red]程序异常退出: {str(e)}[/bold red]")
        if log_level == 'DEBUG':
            logger.exception("详细错误堆栈")
        sys.exit(1)
    finally:
        sys.exit(0)



if __name__ == '__main__':
    run()